# -*- coding: utf-8 -*-
"""
Created on Sun Sep 18 10:45:55 2022
@author: Runx
Describe: This is the Taichi version fo strom_gen.py
"""
import taichi as ti
import settings
import Stochastic_Rainfall_v3 as sr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sys
import math
#sys.tracebacklimit=0

''' Constant '''
resolution = settings.resolution # size of the grid.
ws = settings.ws #storm envelope width.
ls = settings.ls #storm envelope length.
#max_num= settings.max_num 
max_num= sr.max_num_gen() 
grid_n = ws//resolution
grid_m = ls//resolution

''' initialization '''
# num of storm cells
_num = sr.num_gen()
#_pos = sr.loc_gen(_num) # storm cells, row idx and col idx.
#_rmax = sr.r_gen(_num)
#_s1 = sr.s1_gen(_num)
#_s2 = sr.s1_gen(_num)

'''Taichi parallel computing'''
ti.init(arch=ti.cpu)
R = ti.field(float, shape=(grid_m, grid_n)) # Rainfall intensity
P = ti.Vector.field(2, float, shape=(grid_m, grid_n)) # Storm cell position
P_idx = ti.Vector.field(2, float, shape=(grid_m, grid_n)) # convert position itno lattice index 

num = ti.field(int, shape=())
num = _num
s1 = ti.field(float, shape=max_num)
s2 = ti.field(float, shape=max_num)
rmax = ti.field(float, shape=max_num)
loc = ti.Vector.field(2, int, shape=max_num) #location of storm cells.

'''Parameter of Region of Interest (Constant)'''
_roi = settings.roi # a 2D grid with each lattice indexed
dt = settings.dt # unit min
max_dur = settings.max_dur
m, n, tmp = _roi.shape
roi = ti.Vector.field(2, int, shape=(m, n))
roi_pos = ti.Vector.field(2, float, shape=(m, n))
roi.from_numpy(_roi)
roi_pos.from_numpy(_roi)

roi_R = ti.field(float, shape=(m, n))
theta = 30/180 * math.pi # direction of envelope motivation, radian
dur = int(sr.duration_gen()*60/dt) #Duration, Converted to steps.
vel = sr.velocity_gen() # Velocity, (km/hr)
v_1, v_2 = (math.sin(theta), math.cos(theta)) #x,y方向上的速度，km/hr
V = ti.Vector([v_1*1000/60*dt/resolution, v_2*1000/60*dt/resolution])*vel
Rot = ti.Matrix([[0, v_1], [0, v_2]])

def pad_or_cut(value: np.ndarray, target_length: int):
    """pad or cut the rmax, s to the max_num"""
    data_row = None
    if len(value) < target_length:  # 填充
        if len(value.shape)==1:
            data_row = np.pad(value, [(0, target_length - len(value))])
        elif len(value.shape)==2:
            data_row = np.pad(value,[(0, target_length - len(value)), (0, 0)])
    elif len(value) > target_length:  # 截断
        data_row = value[:target_length]
    return data_row

@ti.kernel
def initialize_R():
    '''initialization of R'''
    for i,j in R:
        R[i,j] = 0

@ti.kernel
def initialize_P():
    '''initialization of P'''
    for i,j in P:
        P[i,j] = (i,j)

@ti.func
def neighbor_search(x, y, s1, s2, r_int):
    '''
    Storm cells is an ellipse.
    x, y: the grid idx (position) of the point. (-) 
    s1, s2: the spread of the ellipse. (km)
    '''
    dist1 = ti.math.round(s1*1000/resolution-1, ti.i32) #长半轴
    dist2 = ti.math.round(s2*1000/resolution-1, ti.i32) #短半轴
    x_start = max(x-dist1, 0) 
    y_start = max(y-dist2, 0)
    x_end = min(x + dist1, grid_m)
    y_end = min(y + dist2, grid_n)
    R[x, y]+=r_int
    for i in range(x_start, x_end):
        for j in range(y_start, y_end):
            ellipse = ((i-x)*resolution/1000)**2/s1**2+\
                      ((j-y)*resolution/1000)**2/s2**2
            if ellipse > 1:
                pass
            elif i!=x or j!=y:
                d1 = (i-x)**2
                d2 = (j-y)**2
                R[i,j] += storm_spread(d1, d2, dist1, dist2, r_int)

@ti.func
def storm_spread(d1, d2, dist1, dist2, r_int):
    '''
    rainfall cells are described by a Gaussian function.
    idx : index of the grid in the storm cell, (x, y).
    d1: the distance from the point to Y axis.
    d2: the distance from the point to X axis.
    dist1: major semi axis.
    dist2: minor semi axis.
    '''
    tmp = -d1/(2*dist1**2)\
          -d2/(2*dist2**2)
    r = r_int*ti.math.exp(tmp)
    return(r)

@ti.kernel
def draw_storm():
    for i in loc:
        x0 = loc[i][0]
        y0 = loc[i][1]
        s1_0 = s1[i]
        s2_0 = s2[i]
        r_int = rmax[i]
        #print("每一步的loc不变", x0)
        if i ==0:
            print("Is loc changed?", loc[i][0])
        neighbor_search(x0, y0, s1_0, s2_0, r_int)


@ti.kernel
def roi_trans(x_offset:int, y_offset:int):
    '''
    Put the ROI to the right Position.
    '''
    for i,j in roi:
            roi[i,j] += (x_offset, y_offset)
            roi_pos[i,j] += roi[i,j]

@ti.kernel
def roi_rot():
    '''
    Rotate the Roi idx.
    '''
    for i,j in roi:
        roi[i,j] = Rot @ roi[i,j]
        roi_pos[i,j] += roi[i,j]
        
@ti.kernel
def roi_update_pos():
    for i, j in roi:
        roi_pos[i,j][0] += -vel*1000/60*dt/resolution #只在长轴方向有速度
        roi[i,j][0] = ti.math.round(roi_pos[i,j][0])

@ti.kernel
def roi_record_R():
    for i,j in roi:
        loc = roi[i,j]
        #if 0<=i<grid_m and 0<=j<grid_n:
        if loc[0]>=0:
            roi_R[i,j] += R[loc]

res = []
envelope_R = []

def substeps():
    #initialize_P()
    # rotation firstly, and then tanslation
    global res
    global envelope_R
    roi_rot()
    roi_trans(int(grid_m), int(grid_n/2))
    print(roi)
    for step in range(dur):
        print(step)
        _num = sr.num_gen()
        _pos = sr.loc_gen(_num) # storm cells, row idx and col idx.
        _rmax = sr.r_gen(_num)
        _s1 = sr.s1_gen(_num)
        _s2 = sr.s1_gen(_num)

        # pad or cut. 将暴雨点补充至max_num.
        _pos = pad_or_cut(_pos.values, max_num)
        _s1 = pad_or_cut(_s1, max_num)
        _s2 = pad_or_cut(_s2, max_num)
        _rmax = pad_or_cut(_rmax, max_num)

        s1.from_numpy(_s1)
        s2.from_numpy(_s2)
        rmax.from_numpy(_rmax)
        loc.from_numpy(_pos)
        initialize_R()
        draw_storm()
        roi_update_pos()
        #print(roi)
        roi_record_R()
        res.append(roi_R.to_numpy())
        roi_np = roi.to_numpy()
        print(roi_np)
        max_roi_x = roi_np[:, :, 0].max()
        if max_roi_x <= 0:
            break
        #envelope_R.append(R.to_numpy())

    #print(res)
    print(dur, len(res))
    #np.save("roi.npy",roi.to_numpy())
    #np.save("roi_R.npy", res)
    #np.save("envelope_R.npy", envelope_R)
        #print(loc[0])
        #if step % 5 == 0:
        #    #res = R.to_numpy()
        #    #np.savetxt(f"{step}_R_ti.txt", res)
        #    gui.set_image(roi_R)
        #    gui.show()
    return(res)
gui = ti.GUI('Stochastic Storm', (len(_roi), len(_roi)))
substeps()



    
